from gym.envs.registration import register
import mpe.scenarios as scenarios

def _register(scenario_name, gymkey):
    scenario = scenarios.load(scenario_name + ".py").Scenario()
    world = scenario.make_world()
    name, file_idx, reward_idx, version = gymkey.split('-')
    print(scenario_name, scenario)
    reward_idx2function = {
        '0': scenario.llm0,
        '1': scenario.llm1,
        '2': scenario.llm2,
        '3': scenario.llm3,
        '4': scenario.llm4,
        '5': scenario.llm5,
        '6': scenario.llm6,
    }
    kwargs={
        "world": world,
        "reset_callback": scenario.reset_world,
        "reward_callback": scenario.reward,
        "observation_callback": scenario.observation,
        "get_dead_prey_callback": scenario.get_dead_prey,
        # "info_callback": scenario.info,
        "info_callback": reward_idx2function[reward_idx],
        "done_callback": scenario.done,
    }
    register(
        gymkey,
        entry_point="mpe.environment:MultiAgentEnv",
        kwargs=kwargs,
    )

scenario_name = f"simple_tag-1" # .py name
for j in range(7):
    gymkey = f"SimpleTag-1-{j}-v0" # .yaml name
    _register(scenario_name, gymkey)